#!/usr/bin/env python3

"""
======== DE-CAPTCHA module ========
As `CAPTCHA` is an acronym for "Completely Automated Public Turing test
to tell Computers and Humans Apart".
"""


# built-ins
import pickle
from operator import itemgetter

# site-packages
import numpy as np
import matplotlib.pylab as plt
import matplotlib.cm as cm

from scipy import misc
from scipy import ndimage
from sklearn.cluster import k_means

# local-packages



# Mode Transform

def rgb2gray(im):
    r, g, b = im[:,:,0], im[:,:,1], im[:,:,2]
    return np.array(np.round(0.2989 * r + 0.5870 * g + 0.1140 * b),
                    dtype=np.uint8)

def otsu (im, bins=256):
    steplength = (im.max() - im.min()) / float(bins)
    initial_threshold = im.min() + steplength
    best_bcv = 0
    best_threshold = initial_threshold
    for threshold in np.arange(initial_threshold, im.max(), steplength):
        mask_fg = (im >= threshold)
        mask_bg = (im < threshold)
        wfg = np.count_nonzero(mask_fg)
        wbg = np.count_nonzero(mask_bg)
        if 0 == wfg or 0 == wbg:
            continue
        mfg = im[mask_fg].mean()
        mbg = im[mask_bg].mean()
        bcv = wfg * wbg * (mbg - mfg) ** 2
        if bcv > best_bcv:
            best_bcv = bcv
            best_threshold = threshold
    return best_threshold

def gray2bin(im):
    """Pickout the deep colors."""
    return im < otsu(im, 256)

def rgb2bin(im):
    return gray2bin(rgb2gray(im))


# Clean

def clean(im, structure=np.ones((3, 2))):
    """Using erosion algorithm to cleanup the noises."""
    return ndimage.binary_erosion(im, structure=structure).astype(im.dtype)


# Cluster (Using K-Means)

def cluster(im, k=5):
    h, w = im.shape
    coords = np.array([[j, i] for j in range(h) for i in range(w) if im[j, i]])
    centroids = np.round(k_means(coords, k)[0]).astype(np.int)
    # plt.plot(centroids[1], centroids[0], 'bo')
    return sorted([c for c in centroids], key=itemgetter(1))


# Pick Segment (Using Masks)

def pick_rectangle(im, centroid, half_a, half_b):
    """``a`` for the long side and ``b`` for the short."""
    y, x = centroid
    return im[max(0, y - half_b) : y + half_b, max(0, x - half_a) : x + half_a]

def pick_square(im, centroid, half_a):
    return pick_rectangle(im, centroid, half_a, half_a)

def pick_circle(im, centroid, min_radius=8, max_radius=16, threshold=2):
    h, w = im.shape
    y, x = np.indices(im.shape)
    last = 0
    for radius in range(min_radius, max_radius):
        mask = (y - centroid[0]) ** 2 + (x - centroid[1]) ** 2 < radius ** 2
        select = np.logical_and(mask, im)
        count = np.count_nonzero(select)
        delta = count - last
        if delta <= threshold:
            break
        else:
            last = count
    masked = np.logical_and(mask, im)
    return radius, pick_square(masked, centroid, radius)


# Generalization

def segmentation(im, centroid, delta):
    return pick_circle(im, centroid, 8, 16, delta)[1]

def rotate(im, angle):
    return ndimage.rotate(im, angle, reshape=False, order=0)

def trim(im, crop=1):
    """Trim the ``im`` ignoring the ``crop`` pixels in a line."""
    h, w = im.shape
    x = []
    y = []
    for j in range(h):
        if np.count_nonzero(im[j, :]) > crop:
            y.append(j)
    for i in range(w):
        if np.count_nonzero(im[:, i]) > crop:
            x.append(i)
    left = min(x) if x else 0
    right = (max(x) + 1) if x else w
    top = min(y) if y else 0
    bottom = (max(y) + 1) if y else h
    return im[top:bottom, left:right]

def zoom(im, shape):
    """Zoom the ``im`` to the specific ``shape``."""
    zh, zw = shape
    ih, iw = im.shape
    return ndimage.zoom(im, (zh / ih, zw / iw), order=0)

def generalize(im, shape=(16, 16), angle=30, crop=1):
    """
    Unified function for further processing of the image segments.

    :type im: numpy.ndarray
    :param im: image content array, better to be a bin-value image

    :type shape: tuple/list of len 2 of int
    :param shape: final shape of the output image

    :type angle: int
    :param angle: the range of angles to rotate the image in case of
            finding the thinnest state

    :type crop: int
    :param crop: the max width of lines to ignore when trim the image
    """
    thin = im
    for ang in range(-angle, angle + 1):
        r = rotate(im, ang)
        t = trim(r, crop)
        h, w = t.shape
        if w < thin.shape[1]:
            thin = t
    return zoom(thin, shape)


# Unified Interfaces

def show(im):
    plt.imshow(im, cmap=cm.binary, interpolation='nearest')
    plt.axis('off')
    plt.show()

def similarity(im1, im2):
    v1 = np.packbits(np.uint8(im1)).astype(np.int)
    v2 = np.packbits(np.uint8(im2)).astype(np.int)
    l1 = np.sqrt(v1.dot(v1))
    l2 = np.sqrt(v2.dot(v2))
    return v1.dot(v2) / (l1 * l2)

def decaptcha(im, steady=np.ones((3, 2)), n=5):
    bmp = rgb2bin(im)
    bmp = clean(bmp, structure=steady)
    centroids = cluster(bmp, n)
    segs = [segmentation(bmp, cent, 2) for cent in centroids]
    segs = [generalize(s) for s in segs]
    #
    # recognition code
    #
    return segs


# Once-used Utils

def split_and_dump():
    with open('result.txt') as fp:
        result = [line.strip().split()[:2] for line in fp.readlines()]

    with open('chars.pkl', 'wb') as pkl:
        data = []
        for i, x in result:
            if len(x) != 5:
                continue
            f = 'captcha/%03d.png' % int(i)
            im = misc.imread(f)
            segs = decaptcha(im)
            for j, s in enumerate(segs):
                data.append([s, x[j]])
        pickle.dump(data, pkl)



if __name__ == '__main__':
    pass

    # # im = misc.imread('captcha/018.png')
    # # im = misc.imread('captcha/053.png')
    # gray = rgb2gray(im)
    # bmp = gray2bin(gray)
    # bmp = clean(bmp, structure=np.ones((3, 2)))
    # centroids = cluster(bmp)
    # r, p = pick_circle(bmp, centroids[1])
    # # pt = trim(p)
